import time
import os
from pathlib import Path
import pickle
import numpy as np
import cv2
import open3d as o3d
from semantic_nav.topomap_utils import TopoMap, show_graph_with_images
from semantic_nav.slam_utils import draw_geometries, numpy_to_o3d_pcd, opencv_to_open3d
from semantic_nav.log_utils import get_logger
logger = get_logger()
package_name = 'semantic_nav'
show_framewise_observations = True


def show_topomap(topomap_dir):
    topomap = TopoMap(topomap_dir)
    topomap.show_with_images()

    if show_framewise_observations:
        sorted_observations = sorted(topomap.load_observations().items(), key=lambda x: x[0])
        # Display point clouds observed by each node
        for i, (node, observation) in enumerate(sorted_observations):
            rgbs = [obs[0] for obs in observation]
            rgb = np.hstack(rgbs)
            image = o3d.geometry.Image(cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB))
            point_cloud = observation[0][2]
            # Draw point cloud and image
            draw_geometries([point_cloud, image])

    for i in range(topomap.number_of_nodes()):
        # graph.nodes[i]['pose']: Node pose (R, t) tuple of rotation matrix and translation vector
        # graph.nodes[i]['image_idx']: Image sequence number corresponding to the node
        logger.info(f"node {i} with frame_idx {topomap.node(i)['frame_idx']} and pose {topomap.node(i)['pose']}")

    for u, v, attrs in topomap.edges(data=True):
        # Each edge stores the pose (R, t) tuple of rotation matrix and translation vector and distance relative to the previous node
        logger.info(f"Edge from {u} to {v} with attributes {attrs}")


def main(args=None):
    from ament_index_python.packages import get_package_prefix
    topomap_dir = os.path.join(get_package_prefix(package_name), 'maps')
    show_topomap(topomap_dir)


if __name__ == "__main__":
    main()
